from typing import List

import sentence_transformers
from langchain import FAISS
from langchain.embeddings import HuggingFaceEmbeddings
from chatglm_loader import ChatGLM
# from 压铸机5.MyFAISS import MyFAISS

embedding_model_dict = {\
    "text2vec": "C:\\Work\\llm\\text2vec-large-chinese",
    "m3e-base": "C:\\Work\\llm\\m3e-base"
}

EMBEDDING_MODEL = "m3e-base"
# 初始化 hugginFace 的 embeddings 对象
embeddings = HuggingFaceEmbeddings(model_name=embedding_model_dict[EMBEDDING_MODEL])
embeddings.client = sentence_transformers.SentenceTransformer(
    embeddings.model_name, device='cpu')

db = FAISS.load_local("./faiss/casting2", embeddings)
# 匹配后单段上下文长度
db.chunk_size = 250
db.chunk_conent = True
# 知识检索内容相关度 Score, 数值范围约为0-1100，如果为0，则不生效，经测试设置为小于500时，匹配结果更精准
db.score_threshold = 500

# 选择模型
model = ChatGLM()
model.load_model("C:\\Work\\llm\\ChatGLM2-6B\\THUDM\\chatglm2-6b")

# 基于上下文的prompt模版，请务必保留"{question}"和"{context}"
PROMPT_TEMPLATE = """已知信息：
{context} 

根据上述已知信息，简洁和专业的来回答用户的问题。如果无法从中得到答案，请说 “根据已知信息无法回答该问题” 或 “没有提供足够的相关信息”，不允许在答案中添加编造成分，答案请使用中文。 问题是：{question}"""
def generate_prompt(related_docs: List[str],
                    query: str,
                    prompt_template: str = PROMPT_TEMPLATE, ) -> str:
    context = "\n".join([doc.page_content for doc in related_docs])
    prompt = prompt_template.replace("{question}", query).replace("{context}", context)
    return prompt


# 进行问答
query = "东芝压铸机故障部位压铸机压射前进动作缓慢的故障原因"
related_docs_with_score = db.similarity_search(query, include_metadata=True,k=4)
# related_docs_with_score = db.similarity_search_with_score(query, k=5)
if len(related_docs_with_score) > 0:
    prompt = generate_prompt(related_docs_with_score, query)
else:
    prompt = query


current_length = 0
for answer_result in model.generatorAnswer(prompt=prompt, history=[],
                                                  streaming=True):
        resp = answer_result.llm_output["answer"]
        history = answer_result.history
        history[-1][0] = query
        print(resp[current_length:], end="", flush=True)
        current_length = len(resp)
#       response = {"query": query,
#             "result": resp,
#             "source_documents": related_docs_with_score}

# query = "请用中文回答：东芝压铸机故障部位压铸机安全门报警的故障原因"
# print(qa.run(query))
#
#
# query = "请用中文回答：东芝压铸机故障部位压铸机模具抽芯无动作的故障原因"
# print(qa.run(query))
#
#
# query = "请用中文回答：东芝压铸机故障部位压铸机机床压力异常不增压的故障原因"
# print(qa.run(query))
#
#
# query = "请用中文回答：东芝压铸机故障部位压铸机压射前进动作缓慢的故障原因"
# print(qa.run(query))
